package main

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"errors"
	"flag"
	"fmt"
	"log"
	"math/big"
	"os"
	"strings"
)

const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

func main() {
	envPath := flag.String("env", ".env", "path to the .env file that should be updated")
	length := flag.Int("length", 48, "length of the generated JWT secret")
	flag.Parse()

	secret, err := generateSecret(*length)
	if err != nil {
		log.Fatalf("generate secret: %v", err)
	}

	if err := upsertSecret(*envPath, secret); err != nil {
		log.Fatalf("update %s: %v", *envPath, err)
	}

	fmt.Printf("JWT_SECRET updated in %s\n", *envPath)
}

func generateSecret(length int) (string, error) {
	if length <= 0 {
		return "", fmt.Errorf("length must be greater than zero")
	}

	var builder strings.Builder
	builder.Grow(length)

	for i := 0; i < length; i++ {
		idx, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
		if err != nil {
			return "", err
		}
		builder.WriteByte(charset[idx.Int64()])
	}

	return builder.String(), nil
}

func upsertSecret(path, secret string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			content := fmt.Sprintf("JWT_SECRET=%s\n", secret)
			return os.WriteFile(path, []byte(content), 0o600)
		}
		return err
	}

	var lines []string
	scanner := bufio.NewScanner(bytes.NewReader(data))
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return err
	}

	replaced := false
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.HasPrefix(trimmed, "#") {
			continue
		}
		if idx := strings.Index(line, "="); idx >= 0 {
			key := strings.TrimSpace(line[:idx])
			if key == "JWT_SECRET" {
				lines[i] = fmt.Sprintf("JWT_SECRET=%s", secret)
				replaced = true
				break
			}
		}
	}

	if !replaced {
		lines = append(lines, fmt.Sprintf("JWT_SECRET=%s", secret))
	}

	var builder strings.Builder
	for i, line := range lines {
		if i > 0 {
			builder.WriteByte('\n')
		}
		builder.WriteString(line)
	}
	builder.WriteByte('\n')

	return os.WriteFile(path, []byte(builder.String()), 0o600)
}
